import torch
import torch.nn as nn
from options import HiDDenConfiguration
from model.conv_bn_relu import ConvBNRelu


class Encoder(nn.Module):
    """
    Inserts a watermark into an image.
    """
    def __init__(self, config: HiDDenConfiguration):
        super(Encoder, self).__init__()
        self.H = config.H
        self.W = config.W
        self.conv_channels = config.encoder_channels
        self.num_blocks = config.encoder_blocks

        layers = [ConvBNRelu(3, self.conv_channels)]

        for _ in range(config.encoder_blocks - 1):
            layer = ConvBNRelu(self.conv_channels, self.conv_channels)
            layers.append(layer)

        self.conv_layers = nn.Sequential(*layers)

        # 这里将 config.message_length 改为 1
        self.after_concat_layer = ConvBNRelu(self.conv_channels + 3 + 1, self.conv_channels)
        self.final_layer = nn.Conv2d(self.conv_channels, 3, kernel_size=1)

    def forward(self, image, message):

        # 归一化 message，确保数值范围匹配 image
        message = message.float()  # 如果 image 是 [0,1]，这就够了
        # message = message * 2 - 1  # 如果 image 是 [-1,1]，用这个

        encoded_image = self.conv_layers(image)
        # concatenate message, encoded image, and original image
        concat = torch.cat([message, encoded_image, image], dim=1)
        im_w = self.after_concat_layer(concat)
        im_w = self.final_layer(im_w)
        return im_w
